import numpy as np
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
"""
操作步骤：这一步是准备工作，导入所需的库。
numpy 用于处理数值计算和数组操作， 
torch 是 PyTorch 深度学习框架的核心库，
torch.nn 提供了构建神经网络所需的各种模块，
matplotlib.pyplot 用于数据可视化。
"""

# 造样本数据
w = 2
b = 1
n_features = 1
n_outputs = 1
x = np.linspace(-5.0, 5.0, 100)  # 100 个 x 数据
x = x.reshape(-1, n_features)
y = w * x + b  # 100 个 y 数据
# 将数据转换为 torch.Tensor 类型，使用 torch.float64 以匹配模型参数类型
x = torch.tensor(x, dtype=torch.float64)
y = torch.tensor(y, dtype=torch.float64)
"""
操作步骤：数据准备阶段。
w = 2 和 b = 1 是真实的权重和偏置，用于生成样本数据。
n_features = 1 和 n_outputs = 1 分别表示输入特征的数量和输出的数量，这里是单变量线性回归，
    所以输入和输出特征数都为 1。
np.linspace(-5.0, 5.0, 100) 生成了 100 个在 -5.0 到 5.0 之间均匀分布的 x 数据。
x.reshape(-1, n_features) 将 x 数组重塑为二维数组，以满足 PyTorch 模型输入的要求。
y = w * x + b 根据线性方程生成对应的 y 数据。
torch.tensor() 将 NumPy 数组转换为 PyTorch 张量，并指定数据类型为 torch.float64。
"""

# 定义网络模型
class Net(nn.Module):
    def __init__(self, n_features, n_outputs):
        super(Net, self).__init__()
        # 定义连接神经网络
        self.output = nn.Linear(n_features, n_outputs)  # n_features 输入特征数，n_outputs 输出特征数

    def forward(self, x):
        # 正向传播输入，模型计算输出值
        out = self.output(x)
        return out

net = Net(n_features, n_outputs)  # 实例化网络模型
# 将模型参数转换为 torch.float64 类型
net = net.double()
"""
操作步骤：模型定义阶段。
class Net(nn.Module) 定义了一个继承自 nn.Module 的自定义神经网络类 Net。
__init__ 方法是类的构造函数， super(Net, self).__init__() 调用父类的构造函数。
self.output = nn.Linear(n_features, n_outputs) 定义了一个线性层，该层将输入特征映射到输出特征。
forward 方法定义了模型的正向传播过程，输入 x 通过线性层得到输出 out。
net = Net(n_features, n_outputs) 实例化了 Net 类的对象 net。
net = net.double() 将模型的参数数据类型转换为 torch.float64 ， 以与输入数据类型匹配。
"""

print("w=%.5f" % net.output.weight.item(),
      ",b=%.5f" % net.output.bias.item())
"""
操作步骤：在训练开始前，打印模型的初始权重和偏置，方便后续观察训练过程中参数的变化。
"""

# optmizer 是训练工具
optimizer = torch.optim.SGD(net.parameters(), lr=0.01)  # 优化器，lr 是学习率
loss_func = nn.MSELoss()  # 损失函数，MSELoss 是均方误差损失函数
"""
操作步骤：训练准备阶段。
torch.optim.SGD(net.parameters(), lr=0.01) 定义了随机梯度下降 （ SGD ） 优化器， 用于更新模型的参数。
net.parameters() 表示模型的可训练参数， lr=0.01 是学习率，控制参数更新的步长。
nn.MSELoss() 定义了均方误差 （ MSE ） 损失函数，用于衡量模型预测值与真实值之间的误差。
"""


# 训练模型
for t in range(500):
    o = net(x)
    loss = loss_func(o, y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    # 每 20 个 epoch 打印一次结果
    if t % 20 == 0:
        print("w=%.5f" % net.output.weight.item(),
              ",b=%.5f" % net.output.bias.item(),
              ",loss=%.5f" % loss.item())
"""
操作步骤：模型训练阶段。
for t in range(500) 定义了训练的迭代次数，这里进行 500 次迭代。
o = net(x) 进行正向传播，将输入数据 x 传入模型，得到模型的预测值 o。
loss = loss_func(o, y) 计算预测值 o 与真实值 y 之间的损失。
optimizer.zero_grad() 清空优化器中梯度的缓存，避免梯度累积。
loss.backward() 进行反向传播，计算损失函数关于模型参数的梯度。
optimizer.step() 根据计算得到的梯度，使用优化器更新模型的参数。
if t % 20 == 0 每 20 次迭代打印一次当前的权重、偏置和损失值，方便观察训练过程。
"""

# 绘制样本数据散点图，根据 y 值进行颜色映射
plt.scatter(x.numpy(), y.numpy(), c=y.numpy().flatten(), cmap='viridis')
plt.colorbar(label='y values')
plt.xlabel('x')
plt.ylabel('y')
plt.title('Sample Data Scatter Plot')
plt.show()
"""
操作步骤：数据可视化阶段。
plt.scatter(x.numpy(), y.numpy(), c=y.numpy().flatten(), cmap='viridis') 绘制样本数据的散点图， 
    x.numpy() 和 y.numpy() 将 PyTorch 张量转换为 NumPy 数组， 
    c=y.numpy().flatten() 根据 y 值进行颜色映射，
    cmap='viridis' 指定颜色映射的方案。
plt.colorbar(label='y values') 添加颜色条，并设置标签为 y values。
plt.xlabel('x') 和 plt.ylabel('y') 分别设置 x 轴和 y 轴的标签。
plt.title('Sample Data Scatter Plot') 设置图的标题。
plt.show() 显示绘制的图形。
"""